# 开发者：郭同源
# 开发时间：2021/10/29 22:52
import cv2
import torch,torchvision
import numpy as np
from torch.utils.data import Dataset
import os



class MyData(Dataset):
    def __init__(self,root):
        self.data_addr_list = []
        for filename in os.listdir(root):
            addr = os.path.join(root,filename)
            self.data_addr_list.append(addr)



    def __len__(self):
        return len(self.data_addr_list)


    # 获取图片及相应的信息
    def __getitem__(self, index):
        img_addr = self.data_addr_list[index]
        img = cv2.imread(img_addr)/255#将图像归一化
        img = torch.tensor(img).permute((2,0,1)) #将图像转换成tensor类型
        data_list = img_addr.split('.')  #将图像的文件名进行分割得到 label,position,sort 的信息
        label = int(data_list[1])
        position = data_list[2:6]
        position = [int(i)/300 for i in position] #归一化位置坐标
        sort = int(data_list[6])-1

        return np.float32(img),np.float32(label),np.float32(position),sort






if __name__ == '__main__':


    data = MyData(r'data/train/train')
    # print(data.data_addr_list[0])
    print(data[1][2]*300)
    # print(data[0][0].shape)

    # print(data[0])



